"""
B+树索引实现

实现B+树的核心操作：插入、查找、删除、范围查询
"""

import os
import sys
import struct
from typing import List, Optional, Tuple, Any, Dict
from dataclasses import dataclass

# 添加当前目录到Python路径
current_dir = os.path.dirname(os.path.abspath(__file__))
if current_dir not in sys.path:
    sys.path.insert(0, current_dir)

try:
    from index_types import *
    from constants import PAGE_SIZE
    from statistics import Logger
except ImportError:
    from .index_types import *
    from .constants import PAGE_SIZE
    from .statistics import Logger


class BTreePage:
    """B+树页面基类"""
    
    def __init__(self, page_id: int, page_type: PageType, parent_id: int = 0):
        self.page_id = page_id
        self.page_type = page_type
        self.parent_id = parent_id
        self.is_dirty = False
        self.key_count = 0
        self.keys = []  # IndexKey列表
        
    def to_bytes(self) -> bytes:
        """序列化页面到字节数组"""
        raise NotImplementedError
    
    @classmethod
    def from_bytes(cls, page_id: int, data: bytes):
        """从字节数组反序列化页面"""
        raise NotImplementedError
    
    def is_full(self) -> bool:
        """检查页面是否已满"""
        return self.key_count >= INDEX_CONSTANTS['BTREE_ORDER']
    
    def is_underflow(self) -> bool:
        """检查页面是否需要合并"""
        return self.key_count < INDEX_CONSTANTS['MIN_KEYS']


class BTreeInternalPage(BTreePage):
    """B+树内部节点页面"""
    
    def __init__(self, page_id: int, parent_id: int = 0):
        super().__init__(page_id, PageType.INDEX_INTERNAL, parent_id)
        self.child_pointers = []  # 子页面ID列表
    
    def to_bytes(self) -> bytes:
        """序列化内部节点"""
        # 页面头部：页面类型(1) + 父页面ID(4) + 键数量(4) + 保留(23) = 32字节
        header = struct.pack('<BIII', 
                           self.page_type.value,
                           self.parent_id,
                           self.key_count,
                           0  # 保留字段
                           ) + b'\x00' * 19
        
        # 键和指针交替存储：指针1 + 键1 + 指针2 + 键2 + ... + 指针n+1
        data = b''
        
        # 第一个指针
        if self.child_pointers:
            data += struct.pack('<I', self.child_pointers[0])
        else:
            data += b'\x00\x00\x00\x00'
        
        # 键和后续指针
        for i in range(self.key_count):
            if i < len(self.keys):
                data += self.keys[i].to_bytes()
            else:
                data += b'\x00' * INDEX_CONSTANTS['KEY_SIZE']
            
            if i + 1 < len(self.child_pointers):
                data += struct.pack('<I', self.child_pointers[i + 1])
            else:
                data += b'\x00\x00\x00\x00'
        
        # 填充到页面大小
        total_size = PAGE_SIZE
        current_size = len(header) + len(data)
        padding = b'\x00' * (total_size - current_size)
        
        return header + data + padding
    
    @classmethod
    def from_bytes(cls, page_id: int, data: bytes):
        """从字节数组反序列化内部节点"""
        # 解析头部
        page_type_val, parent_id, key_count, _ = struct.unpack('<BIII', data[:16])
        
        page = cls(page_id, parent_id)
        page.key_count = key_count
        
        # 解析键和指针
        offset = INDEX_CONSTANTS['PAGE_HEADER_SIZE']
        
        # 第一个指针
        if offset + 4 <= len(data):
            first_pointer = struct.unpack('<I', data[offset:offset+4])[0]
            page.child_pointers.append(first_pointer)
            offset += 4
        
        # 键和后续指针
        for i in range(key_count):
            # 读取键（这里简化处理，假设都是整数键）
            if offset + INDEX_CONSTANTS['KEY_SIZE'] <= len(data):
                key_bytes = data[offset:offset + INDEX_CONSTANTS['KEY_SIZE']]
                # 注意：这里需要知道键的原始类型，暂时假设为整数
                try:
                    key_value = struct.unpack('<Q', key_bytes)[0]
                    key = IndexKey(value=key_value, key_type='int')
                    page.keys.append(key)
                except:
                    break
                offset += INDEX_CONSTANTS['KEY_SIZE']
            
            # 读取指针
            if offset + 4 <= len(data):
                pointer = struct.unpack('<I', data[offset:offset+4])[0]
                page.child_pointers.append(pointer)
                offset += 4
        
        return page
    
    def find_child_index(self, key: IndexKey) -> int:
        """找到插入键应该进入的子节点索引"""
        for i, k in enumerate(self.keys):
            if key < k:
                return i
        return len(self.keys)  # 最后一个子节点
    
    def insert_key(self, key: IndexKey, left_child: int, right_child: int) -> bool:
        """插入键和对应的子指针"""
        if self.is_full():
            return False
        
        # 找到插入位置
        insert_pos = 0
        for i, k in enumerate(self.keys):
            if key < k:
                insert_pos = i
                break
            insert_pos = i + 1
        
        # 插入键
        self.keys.insert(insert_pos, key)
        # 更新右子指针
        if insert_pos + 1 < len(self.child_pointers):
            self.child_pointers[insert_pos + 1] = right_child
        else:
            self.child_pointers.append(right_child)
        
        self.key_count += 1
        self.is_dirty = True
        return True


class BTreeLeafPage(BTreePage):
    """B+树叶子节点页面"""
    
    def __init__(self, page_id: int, parent_id: int = 0):
        super().__init__(page_id, PageType.INDEX_LEAF, parent_id)
        self.records = []  # RecordPointer列表
        self.next_leaf = 0  # 下一个叶子页面ID（用于范围查询）
        self.prev_leaf = 0  # 上一个叶子页面ID
    
    def to_bytes(self) -> bytes:
        """序列化叶子节点"""
        # 页面头部：页面类型(1) + 父页面ID(4) + 键数量(4) + 下一个叶子(4) + 上一个叶子(4) + 保留(15) = 32字节
        header = struct.pack('<BIIIIII', 
                           self.page_type.value,
                           self.parent_id,
                           self.key_count,
                           self.next_leaf,
                           self.prev_leaf,
                           0, 0  # 保留字段
                           ) + b'\x00' * 3
        
        # 键值对存储：键1 + 记录指针1 + 键2 + 记录指针2 + ...
        data = b''
        for i in range(self.key_count):
            if i < len(self.keys):
                data += self.keys[i].to_bytes()
            else:
                data += b'\x00' * INDEX_CONSTANTS['KEY_SIZE']
            
            if i < len(self.records):
                data += self.records[i].to_bytes()
            else:
                data += b'\x00' * INDEX_CONSTANTS['RECORD_ID_SIZE']
        
        # 填充到页面大小
        total_size = PAGE_SIZE
        current_size = len(header) + len(data)
        padding = b'\x00' * (total_size - current_size)
        
        return header + data + padding
    
    @classmethod 
    def from_bytes(cls, page_id: int, data: bytes):
        """从字节数组反序列化叶子节点"""
        # 解析头部
        page_type_val, parent_id, key_count, next_leaf, prev_leaf = struct.unpack('<BIIII', data[:20])
        
        page = cls(page_id, parent_id)
        page.key_count = key_count
        page.next_leaf = next_leaf
        page.prev_leaf = prev_leaf
        
        # 解析键值对
        offset = INDEX_CONSTANTS['PAGE_HEADER_SIZE']
        entry_size = INDEX_CONSTANTS['KEY_SIZE'] + INDEX_CONSTANTS['RECORD_ID_SIZE']
        
        for i in range(key_count):
            if offset + entry_size <= len(data):
                # 读取键
                key_bytes = data[offset:offset + INDEX_CONSTANTS['KEY_SIZE']]
                try:
                    key_value = struct.unpack('<Q', key_bytes)[0]
                    key = IndexKey(value=key_value, key_type='int')
                    page.keys.append(key)
                except:
                    break
                
                # 读取记录指针
                record_bytes = data[offset + INDEX_CONSTANTS['KEY_SIZE']:offset + entry_size]
                record_pointer = RecordPointer.from_bytes(record_bytes)
                page.records.append(record_pointer)
                
                offset += entry_size
        
        return page
    
    def insert_record(self, key: IndexKey, record: RecordPointer) -> bool:
        """插入记录到叶子节点"""
        if self.is_full():
            return False
        
        # 找到插入位置（保持有序）
        insert_pos = 0
        for i, k in enumerate(self.keys):
            if key < k:
                insert_pos = i
                break
            elif key == k:
                # 主键重复
                raise DuplicateKeyException(f"Duplicate key: {key.value}")
            insert_pos = i + 1
        
        # 插入键和记录
        self.keys.insert(insert_pos, key)
        self.records.insert(insert_pos, record)
        self.key_count += 1
        self.is_dirty = True
        return True
    
    def find_record(self, key: IndexKey) -> Optional[RecordPointer]:
        """查找记录"""
        for i, k in enumerate(self.keys):
            if k == key:
                return self.records[i]
        return None
    
    def delete_record(self, key: IndexKey) -> bool:
        """删除记录"""
        for i, k in enumerate(self.keys):
            if k == key:
                self.keys.pop(i)
                self.records.pop(i)
                self.key_count -= 1
                self.is_dirty = True
                return True
        return False


class BTreeIndex:
    """B+树索引主类"""
    
    def __init__(self, metadata: IndexMetadata, page_manager, logger: Logger = None):
        self.metadata = metadata
        self.page_manager = page_manager
        self.logger = logger or Logger()
        self.page_cache = {}  # 页面缓存
        self.dirty_pages = set()  # 脏页集合
        
        # 添加B+树操作锁，确保线程安全
        import threading
        self._btree_lock = threading.RLock()  # 使用递归锁
    
    def _load_page(self, page_id: int) -> BTreePage:
        """加载页面"""
        if page_id in self.page_cache:
            return self.page_cache[page_id]
        
        try:
            # 从页面管理器读取原始数据
            page_data = self.page_manager.read_page(page_id)
            
            # 根据页面类型创建相应的页面对象
            page_type_val = struct.unpack('<B', page_data[:1])[0]
            page_type = PageType(page_type_val)
            
            if page_type == PageType.INDEX_INTERNAL:
                page = BTreeInternalPage.from_bytes(page_id, page_data)
            elif page_type == PageType.INDEX_LEAF:
                page = BTreeLeafPage.from_bytes(page_id, page_data)
            else:
                raise IndexCorruptedException(f"Invalid page type for index page: {page_type}")
            
            self.page_cache[page_id] = page
            return page
            
        except Exception as e:
            self.logger.error("INDEX", page_id, f"Failed to load index page: {e}")
            raise IndexException(f"Failed to load page {page_id}: {e}")
    
    def _save_page(self, page: BTreePage):
        """保存页面"""
        if not page.is_dirty:
            return
        
        try:
            page_data = page.to_bytes()
            self.page_manager.write_page(page.page_id, page_data)
            page.is_dirty = False
            self.dirty_pages.discard(page.page_id)
            
        except Exception as e:
            self.logger.error("INDEX", page.page_id, f"Failed to save index page: {e}")
            raise IndexException(f"Failed to save page {page.page_id}: {e}")
    
    def insert(self, key: Any, record_pointer: RecordPointer):
        """插入键值对到索引"""
        # 创建索引键
        index_key = IndexKey(value=key, key_type=self.metadata.key_type)
        
        # 如果是空树，创建根节点
        if self.metadata.root_page_id == 0:
            self._create_root(index_key, record_pointer)
            return
        
        # 查找叶子节点
        leaf_page = self._find_leaf(index_key)
        
        try:
            # 尝试插入到叶子节点
            leaf_page.insert_record(index_key, record_pointer)
            self.dirty_pages.add(leaf_page.page_id)
            self.metadata.entry_count += 1
            
        except DuplicateKeyException:
            raise  # 重新抛出重复键异常
        except:
            # 叶子节点已满，需要分裂
            self._split_leaf(leaf_page, index_key, record_pointer)
    
    def search(self, key: Any) -> Optional[RecordPointer]:
        """查找键对应的记录指针"""
        if self.metadata.root_page_id == 0:
            return None
        
        index_key = IndexKey(value=key, key_type=self.metadata.key_type)
        leaf_page = self._find_leaf(index_key)
        return leaf_page.find_record(index_key)
    
    def delete(self, key: Any) -> bool:
        """删除键"""
        with self._btree_lock:
            if self.metadata.root_page_id == 0:
                return False
            
            index_key = IndexKey(value=key, key_type=self.metadata.key_type)
            leaf_page = self._find_leaf(index_key)
            
            if leaf_page.delete_record(index_key):
                self.dirty_pages.add(leaf_page.page_id)
                self.metadata.entry_count -= 1
                
                # 检查是否需要合并（简化实现，暂不处理合并）
                return True
            
            return False
    
    def _create_root(self, key: IndexKey, record: RecordPointer):
        """创建根节点（叶子节点）"""
        # 分配新页面作为根节点
        root_page_id = self.page_manager.allocate_page()
        self.metadata.root_page_id = root_page_id
        self.metadata.height = 1
        
        # 创建根叶子节点
        root_page = BTreeLeafPage(root_page_id)
        root_page.insert_record(key, record)
        
        self.page_cache[root_page_id] = root_page
        self.dirty_pages.add(root_page_id)
        self.metadata.entry_count = 1
    
    def _find_leaf(self, key: IndexKey) -> BTreeLeafPage:
        """查找键应该在的叶子节点"""
        current_page_id = self.metadata.root_page_id
        
        # 从根节点开始向下查找
        while True:
            current_page = self._load_page(current_page_id)
            
            if isinstance(current_page, BTreeLeafPage):
                return current_page
            
            # 内部节点，找到下一个子节点
            if isinstance(current_page, BTreeInternalPage):
                child_index = current_page.find_child_index(key)
                if child_index < len(current_page.child_pointers):
                    current_page_id = current_page.child_pointers[child_index]
                else:
                    raise IndexCorruptedException("Invalid child pointer in internal node")
            else:
                raise IndexCorruptedException("Invalid page type in index tree")
    
    def _split_leaf(self, leaf_page: BTreeLeafPage, new_key: IndexKey, new_record: RecordPointer):
        """分裂叶子节点"""
        # 分配新的叶子页面
        new_leaf_id = self.page_manager.allocate_page()
        new_leaf = BTreeLeafPage(new_leaf_id, leaf_page.parent_id)
        
        # 合并所有记录（包括新记录）
        all_keys = leaf_page.keys[:]
        all_records = leaf_page.records[:]
        
        # 找到新记录的插入位置
        insert_pos = len(all_keys)
        for i, key in enumerate(all_keys):
            if new_key < key:
                insert_pos = i
                break
        
        all_keys.insert(insert_pos, new_key)
        all_records.insert(insert_pos, new_record)
        
        # 计算分裂点
        mid_point = len(all_keys) // 2
        
        # 更新原叶子节点
        leaf_page.keys = all_keys[:mid_point]
        leaf_page.records = all_records[:mid_point]
        leaf_page.key_count = len(leaf_page.keys)
        leaf_page.is_dirty = True
        
        # 初始化新叶子节点
        new_leaf.keys = all_keys[mid_point:]
        new_leaf.records = all_records[mid_point:]
        new_leaf.key_count = len(new_leaf.keys)
        new_leaf.is_dirty = True
        
        # 更新叶子节点链表
        new_leaf.next_leaf = leaf_page.next_leaf
        new_leaf.prev_leaf = leaf_page.page_id
        leaf_page.next_leaf = new_leaf.page_id
        
        # 缓存新页面
        self.page_cache[new_leaf_id] = new_leaf
        self.dirty_pages.add(leaf_page.page_id)
        self.dirty_pages.add(new_leaf_id)
        
        # 向父节点插入分裂键
        split_key = new_leaf.keys[0]  # 新叶子节点的第一个键作为分裂键
        self._insert_into_parent(leaf_page, split_key, new_leaf)
    
    def _insert_into_parent(self, left_page: BTreePage, key: IndexKey, right_page: BTreePage):
        """向父节点插入分裂键"""
        if left_page.parent_id == 0:
            # 需要创建新的根节点
            self._create_new_root(left_page, key, right_page)
            return
        
        # 加载父节点
        parent_page = self._load_page(left_page.parent_id)
        if not isinstance(parent_page, BTreeInternalPage):
            raise IndexCorruptedException("Parent page is not an internal page")
        
        # 更新子页面的父指针
        right_page.parent_id = parent_page.page_id
        
        # 尝试插入到父节点
        if parent_page.insert_key(key, left_page.page_id, right_page.page_id):
            self.dirty_pages.add(parent_page.page_id)
        else:
            # 父节点也满了，需要分裂（简化实现，暂不处理）
            raise IndexException("Parent node split not implemented")
    
    def _create_new_root(self, left_page: BTreePage, key: IndexKey, right_page: BTreePage):
        """创建新的根节点"""
        # 分配新页面作为根节点
        new_root_id = self.page_manager.allocate_page()
        new_root = BTreeInternalPage(new_root_id)
        
        # 设置根节点的子指针和键
        new_root.child_pointers = [left_page.page_id, right_page.page_id]
        new_root.keys = [key]
        new_root.key_count = 1
        new_root.is_dirty = True
        
        # 更新子页面的父指针
        left_page.parent_id = new_root_id
        right_page.parent_id = new_root_id
        left_page.is_dirty = True
        right_page.is_dirty = True
        
        # 更新索引元数据
        self.metadata.root_page_id = new_root_id
        self.metadata.height += 1
        
        # 缓存新根节点
        self.page_cache[new_root_id] = new_root
        self.dirty_pages.add(new_root_id)
        self.dirty_pages.add(left_page.page_id)
        self.dirty_pages.add(right_page.page_id)
    
    def flush_all(self):
        """刷新所有脏页"""
        with self._btree_lock:
            for page_id in list(self.dirty_pages):
                if page_id in self.page_cache:
                    self._save_page(self.page_cache[page_id])
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取索引统计信息"""
        return {
            'index_name': self.metadata.index_name,
            'table_name': self.metadata.table_name,
            'column_name': self.metadata.column_name,
            'index_type': self.metadata.index_type.value,
            'entry_count': self.metadata.entry_count,
            'height': self.metadata.height,
            'root_page_id': self.metadata.root_page_id,
            'cached_pages': len(self.page_cache),
            'dirty_pages': len(self.dirty_pages)
        }
